#include <iostream>

#include "Image.H"
#include "Util.H"
#include "IO.H"
#include "NearestMeanClassifier.h"
#include "Sample.h"
#include "DataSet.h"

using namespace std;
using namespace pip;
// simpledigitclassifier -trainImages filename -trainLabels filename -testImages filename  -result filename

void pixel_extract(const Image<unsigned char>& in, vector<unsigned char>& h)
{
	h.clear(); // clear vector to avoid side-effects
	for(int x=0;x<in.dim(0);++x){
		for(int y=0;y<in.dim(1);++y){
			h.push_back(in(x,y));
		}
	}
}

void getFeatureLabels(const string images,
                      DataSet& ds,
					  string labels
                      )
{
  Image<unsigned char> labs;
  Image<unsigned char> data;
  Image<unsigned char> in;
  vector<unsigned char> v;

  if( !importFile(labs, labels)){
	  error("while parsing trainings data: ", "Usage:  -trainLabels FILENAME  is not set proper.");
  }
  if( !importFile(data, images)){
	  error("while parsing trainings data: ", "Usage:  -trainImages FILENAME  is not set proper.");
  }
  unsigned int depth = data.dim(2);
  unsigned int width = data.dim(0);
  
  for (unsigned int i = 0; i < depth; ++i) {
	in = data.sliced(2,i);
    pixel_extract(in, v);
    Sample s(v, labs(i,0));
	if(!(i%250))
	  cout << " " << i << " ";
	else if(!(i%10))
	  cout << ".";
    ds.push_back(s);
  }
}

void getFeature(const string images,
                      DataSet& ds
                      )
{
  Image<unsigned char> data;
  Image<unsigned char> in;
  vector<unsigned char> v;

  if( !importFile(data, images)){
	  error("while parsing trainings data: ", "Usage:  -trainImages FILENAME  is not set proper.");
  }
  unsigned int depth = data.dim(2);
  unsigned int width = data.dim(0);
  for (unsigned int i = 0; i < depth; ++i) {
	in = data.sliced(2,i);
    pixel_extract(in, v);
    Sample s(v, 0);
    ds.push_back(s);
  }
}




int main(int argc, char *argv[])
{
  const string program = argv[0];
  //  simpledigitclassifier -trainImages filename -trainLabels filename -testImages filename  -result filename
  vector<string> args(argv, argv + argc);
  if (argc < 6) {
    error(program, "Usage: PROGRAM -trainImages FILENAME -trainLabels FILENAME -testImages FILENAME  -result FILENAME");
  }

  string trIm(" "), trLa(" "), teIm(" "), result("results.pip");//,reference(" ");

  getScalarArg(args, "-trainImages", trIm); 
  getScalarArg(args, "-trainLabels", trLa); 
  getScalarArg(args, "-testImages" , teIm); 
  getScalarArg(args, "-result"     , result); 
  //getScalarArg(args, "-ref"        , reference); 

 

  //variables and calls for experiments
  vector<string> trainfiles, testfiles;
  //getFiles("train", trainfiles);
  //getFiles("test", testfiles);

  DataSet trainingds, testds;

  cout << "=>reading samples..." << endl;
  getFeatureLabels(trIm, trainingds, trLa);
  cout << "=>training means..." << endl;
  NearestMeanClassifier nm;
  nm.train(trainingds);
  cout << "=>preparing test samples..." << endl;
  getFeature(teIm,  testds);
  
  cout << "=>constructing results..." << endl;
  unsigned int length = testds.samples().size();
  Image<unsigned char> out(Dimension(length,0));

  Image<unsigned char> ref;

  for(unsigned int len = 0 ; len < length ; ++len){
	out(len,0) = nm.test(testds.sample(len));
  }
  
  if(!exportFile(out, result)){
	error(program, "...fails to save result \a...");
  }
  
  cout << "=>DONE" << endl;

  //system("pause");
  return 0;
}